//Modified from opencv 2.1
//by Chen Feng (cforrest at umich.edu)
#include "phonyferns.h"
#include <iostream>
#include <fstream>
#include <string>
using namespace std;

const int progressBarSize = 50;    
/////////////////////////////////////// PhonyFernClassifier ////////////////////////////////////////////    

PhonyFernClassifier::PhonyFernClassifier()
{
	verbose = false;
	clear();
}

PhonyFernClassifier::~PhonyFernClassifier()
{
}

void PhonyFernClassifier::clear()
{
	signatureSize = nclasses = nstructs = structSize = compressionMethod = leavesPerStruct = 0;
	vector<Feature>().swap(features);
	vector<float>().swap(posteriors);
}

int PhonyFernClassifier::getLeaf(int fern, const Mat& _patch) const
{
	assert( 0 <= fern && fern < nstructs );
	size_t fofs = fern*structSize, idx = 0;
	const Mat_<uchar>& patch = (const Mat_<uchar>&)_patch;

	for( int i = 0; i < structSize; i++ )
	{
		const Feature& f = features[fofs + i];
		idx = (idx << 1) + f(patch);
	}

	return fern*leavesPerStruct + idx;
}


void PhonyFernClassifier::prepare(
	int _nclasses,
	int _patchSize,
	int _signatureSize,
	int _nstructs, int _structSize,
	int _nviews, int _compressionMethod
	)
{
	clear();

	CV_Assert( _nclasses > 1 && _patchSize >= 5 && _nstructs > 0 &&
		_nviews > 0 && _structSize > 0 &&
		(_compressionMethod == COMPRESSION_NONE ||
		_compressionMethod == COMPRESSION_RANDOM_PROJ ||
		_compressionMethod == COMPRESSION_PCA) );

	nclasses = _nclasses;
	patchSize = Size(_patchSize, _patchSize);
	nstructs = _nstructs;
	structSize = _structSize;
	signatureSize = std::min(_signatureSize, nclasses);
	compressionMethod = signatureSize == 
		nclasses ? COMPRESSION_NONE : _compressionMethod;

	leavesPerStruct = 1 << structSize;

	int i, nfeatures = structSize*nstructs;

	features = vector<Feature>( nfeatures );
	posteriors = vector<float>( leavesPerStruct*nstructs*nclasses, 1.f );
	classCounters = vector<int>( nclasses, leavesPerStruct );

	CV_Assert( patchSize.width <= 256 && patchSize.height <= 256 );
	RNG& rng = theRNG();

	for( i = 0; i < nfeatures; i++ )
	{
		int x1 = (unsigned)rng % patchSize.width;
		int y1 = (unsigned)rng % patchSize.height;
		int x2 = (unsigned)rng % patchSize.width;
		int y2 = (unsigned)rng % patchSize.height;
		features[i] = Feature(x1, y1, x2, y2);
	}
}

void PhonyFernClassifier::trainFromSingleView(
	const Mat& image,
	const vector<KeyPoint>& keypoints,
	int _patchSize, int _signatureSize,
	int _nstructs, int _structSize,
	int _nviews, int _compressionMethod,
	const PatchGenerator& patchGenerator
	)
{
	prepare((int)keypoints.size(), _patchSize, _signatureSize, _nstructs,
		_structSize, _nviews, _compressionMethod);
	int i, j, k, nsamples = (int)keypoints.size(), maxOctave = 0;
	for( i = 0; i < nsamples; i++ )
	{
		classCounters[i] = _nviews;
		//maxOctave = std::max(maxOctave, keypoints[i].octave);
	}

	double maxScale = patchGenerator.lambdaMax*2;
	Mat canvas(cvRound(std::max(image.cols,image.rows)*maxScale 
				+ patchSize.width*2 + 10),
				cvRound(std::max(image.cols,image.rows)*maxScale 
				+ patchSize.width*2 + 10), image.type());
	Mat noisebuf;
	vector<Mat> pyrbuf(maxOctave+1), pyr(maxOctave+1);
	Point2f center0((image.cols-1)*0.5f, (image.rows-1)*0.5f),
		center1((canvas.cols - 1)*0.5f, (canvas.rows - 1)*0.5f);
	Mat matM(2, 3, CV_64F);
	double *M = (double*)matM.data;
	RNG& rng = theRNG();

	Mat patch(patchSize, CV_8U);

	for( i = 0; i < _nviews; i++ )
	{
		patchGenerator.generateRandomTransform(center0, center1, matM, rng);

		CV_Assert(matM.type() == CV_64F);
		Rect roi(INT_MAX, INT_MAX, INT_MIN, INT_MIN);

		for( k = 0; k < 4; k++ )
		{
			Point2f pt0, pt1;
			pt0.x = (float)(k == 0 || k == 3 ? 0 : image.cols);
			pt0.y = (float)(k < 2 ? 0 : image.rows);
			pt1.x = (float)(M[0]*pt0.x + M[1]*pt0.y + M[2]);
			pt1.y = (float)(M[3]*pt0.x + M[4]*pt0.y + M[5]);

			roi.x = std::min(roi.x, cvFloor(pt1.x));
			roi.y = std::min(roi.y, cvFloor(pt1.y));
			roi.width = std::max(roi.width, cvCeil(pt1.x));
			roi.height = std::max(roi.height, cvCeil(pt1.y));
		}

		roi.width -= roi.x + 1;
		roi.height -= roi.y + 1;

		Mat canvas_roi(canvas, roi);
		M[2] -= roi.x;
		M[5] -= roi.y;

		Size size = canvas_roi.size();
		rng.fill(canvas_roi, RNG::UNIFORM, Scalar::all(0), Scalar::all(256));
		warpAffine( image, canvas_roi, matM, size, INTER_LINEAR, BORDER_TRANSPARENT);

		pyr[0] = canvas_roi;                
		/*for( j = 1; j <= maxOctave; j++ )
		{
			size = Size((size.width+1)/2, (size.height+1)/2);
			if( pyrbuf[j].cols < size.width*size.height )
				pyrbuf[j].create(1, size.width*size.height, image.type());
			pyr[j] = Mat(size, image.type(), pyrbuf[j].data);
			pyrDown(pyr[j-1], pyr[j]);
		}*/

		if( patchGenerator.noiseRange > 0 )
		{
			const int noiseDelta = 128;
			if( noisebuf.cols < pyr[0].cols*pyr[0].rows )
				noisebuf.create(1, pyr[0].cols*pyr[0].rows, image.type());
			for( j = 0; j <= maxOctave; j++ )
			{
				Mat noise(pyr[j].size(), image.type(), noisebuf.data);
				rng.fill(noise, RNG::UNIFORM, Scalar::all(-patchGenerator.noiseRange + noiseDelta),
					Scalar::all(patchGenerator.noiseRange + noiseDelta));
				addWeighted(pyr[j], 1, noise, 1, -noiseDelta, pyr[j]);
			}
		}

		for( j = 0; j < nsamples; j++ )
		{
			KeyPoint kpt = keypoints[j];
			float scale = 1.f/(1 << kpt.octave);
			Point2f pt((float)((M[0]*kpt.pt.x + M[1]*kpt.pt.y + M[2])*scale),
				(float)((M[3]*kpt.pt.x + M[4]*kpt.pt.y + M[5])*scale));
			getRectSubPix(pyr[kpt.octave], patchSize, pt, patch, patch.type());
			for( int f = 0; f < nstructs; f++ )
				posteriors[getLeaf(f, patch)*nclasses + j]++;
		}

		if( verbose && (i+1)*progressBarSize/_nviews != i*progressBarSize/_nviews )
			putchar('.');
	}
	if( verbose )
		putchar('\n');

	finalize(rng);
}


int PhonyFernClassifier::operator()(const Mat& img, Point2f pt, vector<float>& signature) const
{
	Mat patch;
	getRectSubPix(img, patchSize, pt, patch, img.type());
	return (*this)(patch, signature);
}

int PhonyFernClassifier::operator()(const Mat& patch, vector<float>& signature) const
{
	if( posteriors.empty() )
		CV_Error(CV_StsNullPtr,
		"The descriptor has not been trained or "
		"the floating-point posteriors have been deleted");
	CV_Assert(patch.size() == patchSize);

	int i, j, sz = signatureSize;
	signature.resize(sz);
	float* s = &signature[0];

	for( j = 0; j < sz; j++ )
		s[j] = 0;

	for( i = 0; i < nstructs; i++ )
	{
		int lf = getLeaf(i, patch);
		const float* ldata = &posteriors[lf*signatureSize];
		for( j = 0; j <= sz - 4; j += 4 )
		{
			float t0 = s[j] + ldata[j];
			float t1 = s[j+1] + ldata[j+1];
			s[j] = t0; s[j+1] = t1;
			t0 = s[j+2] + ldata[j+2];
			t1 = s[j+3] + ldata[j+3];
			s[j+2] = t0; s[j+3] = t1;
		}
		for( ; j < sz; j++ )
			s[j] += ldata[j];
	}

	j = 0;
	if( signatureSize == nclasses && compressionMethod == COMPRESSION_NONE )
	{
		for( i = 1; i < nclasses; i++ )
			if( s[j] < s[i] )
				j = i;
	}
	return j;
}


void PhonyFernClassifier::finalize(RNG&)
{
	int i, j, k, n = nclasses;
	vector<double> invClassCounters(n);
	Mat_<double> _temp(1, n);
	double* temp = &_temp(0,0);

	for( i = 0; i < n; i++ )
		invClassCounters[i] = 1./classCounters[i];

	for( i = 0; i < nstructs; i++ )
	{
		for( j = 0; j < leavesPerStruct; j++ )
		{
			float* P = &posteriors[(i*leavesPerStruct + j)*nclasses];
			double sum = 0;
			for( k = 0; k < n; k++ )
				sum += P[k]*invClassCounters[k];
			sum = 1./sum;
			for( k = 0; k < n; k++ )
				temp[k] = P[k]*invClassCounters[k]*sum;
			log(_temp, _temp);
			for( k = 0; k < n; k++ )
				P[k] = (float)temp[k];
		}
	}
}

void PhonyFernClassifier::setVerbose(bool _verbose)
{
	verbose = _verbose;
}    

////////////////////////////////////// Planar Object Detector ////////////////////////////////////

PhonyPlanarObjectDetector::PhonyPlanarObjectDetector()
{
}

PhonyPlanarObjectDetector::~PhonyPlanarObjectDetector()
{
}    

vector<KeyPoint> PhonyPlanarObjectDetector::getModelPoints()
{
	return modelPoints;
}   

void PhonyPlanarObjectDetector::train(
	const Mat& pyr,
	const vector<KeyPoint>& keypoints,
	int patchSize, int nstructs, int structSize,
	int nviews,
	const PatchGenerator& patchGenerator
	)
{
	modelROI = Rect(0, 0, pyr.cols, pyr.rows);
	modelPoints.resize(keypoints.size());
	std::copy(keypoints.begin(), keypoints.end(), modelPoints.begin());

	fernClassifier.setVerbose(verbose);
	fernClassifier.trainFromSingleView(pyr, modelPoints, 
		patchSize, (int)modelPoints.size(), nstructs, structSize, nviews,
		PhonyFernClassifier::COMPRESSION_NONE, patchGenerator);
}    


#define MYREAD(in,data) (in).read((char*)&(data), sizeof((data)))
#define MYWRITE(out,data) (out).write((char*)&(data), sizeof((data)))

bool PhonyPlanarObjectDetector::read_binary_detector(string path)
{
	//PhonyPlanarObjectDetector& d=*this;
	std::ifstream in(path.c_str(), std::ios::binary);
	if( in.is_open() ) {
		cout<<"[read_binary_detector] reading begin -->"<<endl;
	} else {
		cout<<"[read_binary_detector] can not open the file :"<<endl;
		cout<<path<<endl;
		cout<<"[read_binary_detector] <-- reading exit!"<<endl;
		return false;
	}
	{//1. ROI
		cout<<"[read_binary_detector] reading ROI..."<<endl;
		MYREAD(in, modelROI.x);
		MYREAD(in, modelROI.y);
		MYREAD(in, modelROI.width);
		MYREAD(in, modelROI.height);
	}
	{//2. model-points
		cout<<"[read_binary_detector] reading model-points..."<<endl;
		int n;
		MYREAD(in, n);
		modelPoints.resize(n);
		for(int i=0; i<n; ++i) {
			KeyPoint& k=modelPoints[i];
			MYREAD(in, k.pt.x);
			MYREAD(in, k.pt.y);
			MYREAD(in, k.size);
			MYREAD(in, k.angle);
			MYREAD(in, k.response);
			MYREAD(in, k.octave);
			MYREAD(in, k.class_id);
		}
	}
	{//3. fern-classifier
		cout<<"[read_binary_detector] reading fern-classifier..."<<endl;
		PhonyFernClassifier& fs = this->fernClassifier;
		fs.setVerbose(true);
		MYREAD(in, fs.nstructs);
		MYREAD(in, fs.structSize);
		MYREAD(in, fs.nclasses);
		MYREAD(in, fs.signatureSize);
		MYREAD(in, fs.compressionMethod);
		MYREAD(in, fs.leavesPerStruct);
		MYREAD(in, fs.patchSize.width);
		MYREAD(in, fs.patchSize.height);
		int nf;
		MYREAD(in,nf);
		fs.features.resize(nf);
		for(int i = 0; i < nf; ++i)
		{
			int ofs1, ofs2;
			MYREAD(in,ofs1);
			MYREAD(in,ofs2);
			fs.features[i] = PhonyFernClassifier::Feature(
				ofs1%fs.patchSize.width, ofs1/fs.patchSize.width,
				ofs2%fs.patchSize.width, ofs2/fs.patchSize.width  );
		}
		{//4. posteriors
			cout<<"[read_binary_detector] reading posteriors..."<<endl;
			int np;
			MYREAD(in,np);
			fs.posteriors.resize(np);
			for(int i=0; i<np; ++i) {
				MYREAD(in, fs.posteriors[i]);
			}
		}
	}
	in.close();
	cout<<"[read_binary_detector] <-- reading done!"<<endl;
	return true;
}

bool PhonyPlanarObjectDetector::write_binary_detector(string path) const
{
	//const PhonyPlanarObjectDetector& d=*this;
	std::ofstream o(path.c_str(), std::ios::binary);
	if( o.is_open() ) {
		cout<<"[write_binary_detector] writing begin -->"<<endl;
	} else {
		cout<<"[write_binary_detector] can not open the file :"<<endl;
		cout<<path<<endl;
		cout<<"[write_binary_detector] <-- writing exit!"<<endl;
		return false;
	}
	{//1. ROI
		cout<<"[write_binary_detector] writing ROI..."<<endl;
		MYWRITE(o, modelROI.x);
		MYWRITE(o, modelROI.y);
		MYWRITE(o, modelROI.width);
		MYWRITE(o, modelROI.height);
	}
	{//2. model-points
		cout<<"[write_binary_detector] writing model-points..."<<endl;
		int n = (int) this->modelPoints.size();
		MYWRITE(o, n);
		for(int i=0; i<n; ++i) {
			const KeyPoint& k=modelPoints[i];
			MYWRITE(o, k.pt.x);
			MYWRITE(o, k.pt.y);
			MYWRITE(o, k.size);
			MYWRITE(o, k.angle);
			MYWRITE(o, k.response);
			MYWRITE(o, k.octave);
			MYWRITE(o, k.class_id);
		}
	}
	{//3. fern-classifier
		cout<<"[write_binary_detector] writing fern-classifier..."<<endl;
		const PhonyFernClassifier& fs = this->fernClassifier;
		MYWRITE(o, fs.nstructs);
		MYWRITE(o, fs.structSize);
		MYWRITE(o, fs.nclasses);
		MYWRITE(o, fs.signatureSize);
		MYWRITE(o, fs.compressionMethod);
		MYWRITE(o, fs.leavesPerStruct);
		MYWRITE(o, fs.patchSize.width);
		MYWRITE(o, fs.patchSize.height);
		int nf = (int)fs.features.size();
		MYWRITE(o,nf);
		for(int i = 0; i < nf; ++i)
		{
			int a = (fs.features[i].y1*fs.patchSize.width + fs.features[i].x1);
			int b = (fs.features[i].y2*fs.patchSize.width + fs.features[i].x2);
			MYWRITE(o, a);
			MYWRITE(o, b);
		}
		{//4. posteriors
			cout<<"[write_binary_detector] writing posteriors..."<<endl;
			int ps = (int)fs.posteriors.size();
			MYWRITE(o, ps);
			for(int i=0; i<ps; ++i) {
				MYWRITE(o, fs.posteriors[i]);
			}
		}
	}
	o.close();
	cout<<"[write_binary_detector] <-- writing done!"<<endl;
	return true;
}

bool PhonyPlanarObjectDetector::operator()(
	const Mat& pyr, const vector<KeyPoint>& keypoints,
	vector<Point2f>& toPt,
	Mat& matH, vector<Point2f>& corners, vector<int>* pairs) const
{
	double t = (double)getTickCount();
	int i, j, m = (int)modelPoints.size(), n = (int)keypoints.size();
	vector<int> bestMatches(m, -1);
	vector<float> maxLogProb(m, -FLT_MAX);
	vector<float> signature;
	vector<Point2f> fromPt;

	for( i = 0; i < n; i++ )
	{
		KeyPoint kpt = keypoints[i];
		//CV_Assert(0 <= kpt.octave && kpt.octave < (int)pyr.size());//assume octave is 0, FAST
		//kpt.pt.x /= (float)(1 << kpt.octave);
		//kpt.pt.y /= (float)(1 << kpt.octave);
		int k = fernClassifier(pyr/*[kpt.octave]*/, kpt.pt, signature);
		if( k >= 0 && (bestMatches[k] < 0 || signature[k] > maxLogProb[k]) )
		{
			maxLogProb[k] = signature[k];
			bestMatches[k] = i;
		}
	}

	if(pairs)
		pairs->resize(0);

	for( i = 0; i < m; i++ ) {
		if( bestMatches[i] >= 0 )
		{
			fromPt.push_back(modelPoints[i].pt);
			toPt.push_back(keypoints[bestMatches[i]].pt);
		}
	}

	if( fromPt.size() < 4 )
		return false;

	double tclass = (double)getTickCount()-t;
	cout<<"recognize="<<( tclass*1000/getTickFrequency() )<<" ";

	vector<uchar> mask;
	//FC trade-off: RANSAC/LMEDS
	matH = findHomography(Mat(fromPt), Mat(toPt), mask, RANSAC, 10);
	double transac = (double)getTickCount()- t - tclass;
	cout<<"ransac="<<( transac*1000/getTickFrequency() )<<" ";
	if( matH.data )
	{
		const Mat_<double>& H = matH;
		corners.resize(4);
		for( i = 0; i < 4; i++ )
		{
			Point2f pt((float)(modelROI.x + (i == 0 || i == 3 ? 0 : modelROI.width)),
				(float)(modelROI.y + (i <= 1 ? 0 : modelROI.height)));
			double w = 1./(H(2,0)*pt.x + H(2,1)*pt.y + H(2,2));
			corners[i] = Point2f((float)((H(0,0)*pt.x + H(0,1)*pt.y + H(0,2))*w),
				(float)((H(1,0)*pt.x + H(1,1)*pt.y + H(1,2))*w));
		}
	}

	if( pairs )
	{
		for( i = j = 0; i < m; i++ )
			if( bestMatches[i] >= 0 && mask[j++] )
			{
				pairs->push_back(i);
				pairs->push_back(bestMatches[i]);
			}
	}

	return matH.data != 0;
}

//FC
bool PhonyPlanarObjectDetector::operator()(const Mat& pyr, const vector<KeyPoint>& keypoints,
		vector<int>* pairs, float thresh) const
{
	double t = (double)getTickCount();
	int i, j, m = (int)modelPoints.size(), n = (int)keypoints.size();
	vector<int> bestMatches(m, -1);
	vector<float> maxLogProb(m, -FLT_MAX);
	vector<float> signature;

	float mi=FLT_MAX,ma=-mi;

	for( i = 0; i < n; i++ )
	{
		KeyPoint kpt = keypoints[i];
		int k = fernClassifier(pyr, kpt.pt, signature);
		if( k >= 0 && (bestMatches[k] < 0 || signature[k] > maxLogProb[k]) )
		{
			maxLogProb[k] = signature[k];
			bestMatches[k] = i;
			mi = min(mi, signature[k]);
			ma = max(ma, signature[k]);
		}
	}

	cout<<"("<<mi<<","<<ma<<") ";

	double tclass = (double)getTickCount()-t;
	cout<<"ferns="<<( tclass*1000/getTickFrequency() )<<" ";

	if( pairs )
	{
		pairs->resize(0);
		for( i = j = 0; i < m; i++ )
			if( bestMatches[i] >= 0 && maxLogProb[i]>thresh )
			{
				pairs->push_back(i);
				pairs->push_back(bestMatches[i]);
			}
	}

	return true;
}

void PhonyPlanarObjectDetector::setVerbose(bool _verbose)
{
	verbose = _verbose;
}
